# FLANN匹配

import cv2
import numpy as np
from matplotlib import pyplot as plt

query_image = cv2.imread("image/1.jpg", cv2.IMREAD_GRAYSCALE)
training_image = cv2.imread("image/2.jpg", cv2.IMREAD_GRAYSCALE)

# create SIFT and detect/compute
sift = cv2.xfeatures2d.SIFT_create()        # cow， 又是版权问题，无法执行。
kp1, des1 = sift.detectAndCompute(query_image, None)
kp2, des2 = sift.detectAndCompute(training_image, None)

# FLANN matcher parameters
FLANN_INDEX_KDTREE = 0
index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
search_params = dict(checks=50)

flann = cv2.FlannBasedMatcher(index_params, search_params)

matches = flann.knnMatch(des1, des2, k=2)

# prepare an empty mask to draw good matches.
matches_mask = [
    [0, 0] for i in range(len(matches))
]

for i, (m, n) in enumerate(matches):
    if m.distance < 0.7 * n.distance:
        matches_mask[i] = [1, 0]

draw_params = dict(
    matchColor=(0, 255, 0),
    singlePointColor=(255, 0, 0),
    matchesMask=matches_mask,
    flags=0
)

result_image = cv2.drawMatchesKnn(query_image, kp1, training_image, kp2, matches, None, **draw_params)

plt.imshow(result_image)
plt.show()
